import albumentations as A
import os.path as osp
import os, shutil
import cv2


def show_two_image(image1, image2, title=None):
    # 同时可视化两个RGB或者mask
    from matplotlib import pyplot as plt
    fig = plt.figure(figsize=(10, 10))
    ax1 = plt.subplot(1, 2, 1)
    ax2 = plt.subplot(1, 2, 2)
    plt.sca(ax1)
    plt.imshow(image1)
    plt.sca(ax2)
    plt.imshow(image2)
    if title:
        plt.title(title)
    plt.show()


def pixel_aug_xml():
    "对xml数据使用像素增强，并且将对应的xml改名复制到输出路径"
    image_dir = r"/Users/edy/Data/DB/label_data/drop_left61_mg/JPEGImages"
    xml_dir = r"/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/Annotations"
    out_dir = r"/Users/edy/Data/yq1-0-show/yq1_0_show_data_mix_dp/out_aug"

    out_img_dir = osp.join(out_dir, 'Images')
    out_ann_dir = osp.join(out_dir, 'Annotations')
    # os.makedirs(out_ann_dir, exist_ok=True)
    # os.makedirs(out_img_dir, exist_ok=True)

    # target_files = ['frame_64.jpg', 'frame_56.jpg', 'frame_169.jpg', 'frame_79.jpg', ]
    # files_path = [osp.join(image_dir, x) for x in target_files]

    aug_name_str = 'RBC'  # 增强方式的名称

    transform = A.Compose([
            # A.Blur(blur_limit=9,p=1),
            # A.MotionBlur(blur_limit=(13,16),p=1),
            # A.GaussNoise(p=1),
            # A.Sharpen(p=1),
            # A.RandomBrightnessContrast(p=1),
            # A.CLAHE(p=1),
            # A.Downscale(p=1),
            A.RandomShadow(p=1),
            # A.RandomRain(p=1),
        ])
    import glob
    for file in glob.glob(osp.join(image_dir,"*.jpg")) :
        img = cv2.imread(file, 1)[:,:,::-1]
        aug_img = transform(image=img)['image']

        show_two_image(img,aug_img)
        # show_two_image(img[:,:,::-1],aug_img[:,:,::-1])
        print('ii')
        # name, suffix = osp.basename(file).split('.')
        # new_name = 'aug_%s_%s.%s' % (name, aug_name_str, suffix)
        # 
        # cv2.imwrite(osp.join(out_img_dir, new_name), aug_img, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
        # shutil.copy(osp.join(xml_dir, name + '.xml'), osp.join(out_ann_dir, new_name.split('.')[0] + '.xml'))

pixel_aug_xml()